# coding: utf-8

from numpy import deg2rad, array, cov, var, sqrt, sum, cos, sin, mean
from robotics.io.pcd import load_point_cloud
from robotics.gui.glet import GridWindow
from robotics.registration.correspondence.oneway import find_correspondences
from robotics.registration.util import transform_point_cloud
from robotics.registration.sample.random import sampler


def rotation_around_y_estimation(deg):
    r = deg2rad(deg)
    return array([
        [cos(r), 0, -sin(r)],
        [0, 1, 0],
        [sin(r), 0, cos(r)]
    ])


window = GridWindow(caption='Viz')

# f1 = '/Users/yfwz100/Documents/notebook/icp/robot/data/office2/13-[1.30,0.00,-1.98]-[0.00,-0.87,0.00].pcd'
# f2 = '/Users/yfwz100/Documents/notebook/icp/robot/data/office2/14-[1.36,0.00,-2.06]-[0.00,-0.70,0.00].pcd'

f1 = '/Users/yfwz100/Documents/notebook/icp/robot/data/office2/1-[0.00,0.00,0.00]-[0.00,0.00,0.00].pcd'
f2 = '/Users/yfwz100/Documents/notebook/icp/robot/data/office2/2-[0.07,0.00,-0.19]-[0.00,-0.52,0.00].pcd'

p1 = sampler(load_point_cloud(f1) * 10)
p2 = sampler(load_point_cloud(f2) * 10)

print(p1.shape)
print(p2.shape)

# p1 = transform_point_cloud(p1, rotation_around_y_estimation(-30), [1.9, 0, 0.7])
p1 = transform_point_cloud(p1, rotation_around_y_estimation(-28.89), [2.26, 0, 0.1])

p1c, p2c = find_correspondences(p1, p2, max_dist=5)
# p2c, p1c = find_correspondences(p2, p1, max_dist=5)

err = mean(sqrt(sum((p1c - p2c) ** 2, axis=1)))
print(err)

# v = var(sqrt(sum((p1c - p2c) ** 2, axis=1)))
# print(v)
#
# p1c, p2c = find_correspondences(p2, p1, max_dist=3*v)

print(p1c.shape)


@window.within_draw_process
def draw(gl):
    global p1c, p2c

    gl.glColor3d(0.5, 0.5, 0.5)
    gl.glLineWidth(2)

    for p, x in zip(p1c, p2c):
        gl.glBegin(gl.GL_LINES)
        gl.glVertex3d(p[0] / 100, 0, p[2] / 100)
        gl.glVertex3d(x[0] / 100, 0, x[2] / 100)
        gl.glEnd()

    gl.glPointSize(2)

    gl.glColor3d(0.1, 0.5, 1)  # blue

    gl.glBegin(gl.GL_POINTS)
    for p in p1c:
        gl.glVertex3d(p[0] / 100, 0, p[2] / 100)
    gl.glEnd()

    gl.glColor3d(0.1, 1, 0.5)  # green

    gl.glBegin(gl.GL_POINTS)
    for p in p2c:
        gl.glVertex3d(p[0] / 100, 0, p[2] / 100)
    gl.glEnd()


window.run()
